import torch.nn as nn
import torch


class AddGeoCoords(nn.Module):

    def __init__(self, with_r=False):
        super().__init__()
        self.with_r = with_r

    def __call__(self, tensor_coords):
        input_tensor = tensor_coords[0]
        coords = tensor_coords[1]
        top_left, bottom_right, width, height, pixel_width, pixel_height = coords['top_left'], coords['bottom_right'], \
                                                                           coords['width'], coords['height'], \
                                                                           coords['pixel_width'], coords['pixel_height']

        top_left_x, top_left_y = top_left
        bottom_right_x, bottom_right_y = bottom_right

        # 我们在数据预处理阶段加Coords，所以batch_size只能等于1
        channels, x_dim, y_dim = input_tensor.size()

        # xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1)
        # yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1).transpose(1, 2)

        # 创建 x 和 y 坐标通道
        xx_channel = (torch.arange(x_dim).float() * pixel_width + top_left_x).repeat(1, y_dim, 1)
        yy_channel = (torch.arange(y_dim).float() * pixel_height + top_left_y).repeat(1, x_dim, 1).transpose(1, 2)

        # xx_channel = xx_channel.float() / (x_dim - 1)
        # yy_channel = yy_channel.float() / (y_dim - 1)
        #
        # xx_channel = xx_channel * 2 - 1
        # yy_channel = yy_channel * 2 - 1

        # 归一化坐标通道
        # xx_channel = (xx_channel - top_left_x) / (x_dim * pixel_width) * 2 - 1
        # yy_channel = (yy_channel - top_left_y) / (y_dim * pixel_height) * 2 - 1
        50000,50000
        
        
        70000,60000
        
        xx_channel = xx_channel.transpose(1, 2)
        yy_channel = yy_channel.transpose(1, 2)

        ret = torch.cat([
            input_tensor,
            xx_channel.type_as(input_tensor),
            yy_channel.type_as(input_tensor)], dim=0)

        if self.with_r:
            rr = torch.sqrt(
                torch.pow(xx_channel.type_as(input_tensor) - 0.5, 2) + torch.pow(yy_channel.type_as(input_tensor) - 0.5,
                                                                                 2))
            ret = torch.cat([ret, rr], dim=1)

        return ret

class AddCoords(nn.Module):

    def __init__(self, with_r=False):
        super().__init__()
        self.with_r = with_r

    def forward(self, input_tensor):
        """
        Args:
            input_tensor: shape(batch, channel, x_dim, y_dim)
        """
        batch_size, _, x_dim, y_dim = input_tensor.size()

        xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1)
        yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1).transpose(1, 2)

        
        
        
        xx_channel = xx_channel.float() / (x_dim - 1)
        yy_channel = yy_channel.float() / (y_dim - 1)

        xx_channel = xx_channel * 2 - 1
        yy_channel = yy_channel * 2 - 1

        xx_channel = xx_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3)
        yy_channel = yy_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3)

        ret = torch.cat([
            input_tensor,
            xx_channel.type_as(input_tensor),
            yy_channel.type_as(input_tensor)], dim=1)

        if self.with_r:
            rr = torch.sqrt(
                torch.pow(xx_channel.type_as(input_tensor) - 0.5, 2) + torch.pow(yy_channel.type_as(input_tensor) - 0.5,
                                                                                 2))
            ret = torch.cat([ret, rr], dim=1)

        return ret

class CoordConv(nn.Module):

    def __init__(self, in_channels, out_channels, with_r=False, **kwargs):
        super().__init__()
        self.addcoords = AddCoords(with_r=with_r)
        in_size = in_channels + 2
        if with_r:
            in_size += 1
        self.conv = nn.Conv2d(in_size, out_channels, **kwargs)

    def forward(self, x):
        ret = self.addcoords(x)
        ret = self.conv(ret)
        return ret
